import { MilvusClient, DataType } from '@zilliz/milvus2-sdk-node';
import { VmoRequest } from 'vmo-request';
import * as readline from 'readline';

// 创建 Milvus 客户端
const milvusClient = new MilvusClient({
  address: 'localhost:19530',
});

// 创建请求实例
const vmoRequest = new VmoRequest({
  baseURL: 'http://localhost:11434',
  timeout: 30000,
}, true);

// 创建 readline 接口
const rl = readline.createInterface({
  input: process.stdin,
  output: process.stdout
});

/**
 * 生成文本嵌入
 */
async function generateEmbedding(text: string): Promise<number[]> {
  try {
    const response = await vmoRequest.request({
      url: '/api/embeddings',
      method: 'post',
      data: {
        model: 'bge-m3',
        prompt: text,
      }
    });

    if (!response?.data?.embedding || !Array.isArray(response.data.embedding)) {
      throw new Error('生成嵌入失败：响应格式错误');
    }

    return response.data.embedding;
  } catch (error) {
    console.error('生成嵌入失败:', error);
    throw error;
  }
}

/**
 * 初始化集合
 */
async function initCollection() {
  try {
    // 检查集合是否存在
    const hasCollection = await milvusClient.hasCollection({
      collection_name: 'text_embeddings',
    });

    if (!hasCollection.value) {
      // 创建集合
      await milvusClient.createCollection({
        collection_name: 'text_embeddings',
        fields: [
          {
            name: 'id',
            data_type: DataType.Int64,
            is_primary_key: true,
            autoID: true,
          },
          {
            name: 'text',
            data_type: DataType.VarChar,
            
            max_length: 65535,
          },
          {
            name: 'embedding',
            data_type: DataType.FloatVector,
            dim: 1024,
          },
        ],
      });

      // 创建索引
      await milvusClient.createIndex({
        collection_name: 'text_embeddings',
        field_name: 'embedding',
        index_type: 'HNSW',
        metric_type: 'COSINE',
        params: { M: 8, efConstruction: 64 },
      });

      console.log('集合创建成功');
    }

    // 加载集合
    await milvusClient.loadCollection({
      collection_name: 'text_embeddings',
    });
    console.log('集合加载成功');
  } catch (error) {
    console.error('初始化集合失败:', error);
    throw error;
  }
}

/**
 * 清空集合
 */
async function clearCollection() {
  try {
    await milvusClient.dropCollection({
      collection_name: 'text_embeddings'
    });
    console.log('集合已清空');
    // 重新初始化集合
    await initCollection();
  } catch (error) {
    console.error('清空集合失败:', error);
  }
}

/**
 * 添加向量数据
 */
async function addVectorData() {
  try {
    rl.question('请输入要添加的文本: ', async (text) => {
      try {
        const embedding = await generateEmbedding(text);
        await milvusClient.insert({
          collection_name: 'text_embeddings',
          fields_data: [{
            text,
            embedding,
          }],
        });
        console.log('向量数据添加成功');
      } catch (error) {
        console.error('添加向量数据失败:', error);
      }
      showMenu();
    });
  } catch (error) {
    console.error('操作失败:', error);
    showMenu();
  }
}

/**
 * 查询向量数据
 */
async function searchVectors() {
  try {
    rl.question('请输入查询文本: ', async (queryText) => {
      try {
        const queryEmbedding = await generateEmbedding(queryText);
        const results = await milvusClient.search({
          collection_name: 'text_embeddings',
          vector: queryEmbedding,
          output_fields: ['text'],
          limit: 5,
          metric_type: 'COSINE',
        });
        console.log('查询结果:', JSON.stringify(results, null, 2));
      } catch (error) {
        console.error('查询失败:', error);
      }
      showMenu();
    });
  } catch (error) {
    console.error('操作失败:', error);
    showMenu();
  }
}

/**
 * 显示菜单
 */
function showMenu() {
  console.log('\n=== Milvus 向量数据库管理工具 ===');
  console.log('1. 清空集合');
  console.log('2. 添加向量数据');
  console.log('3. 查询向量数据');
  console.log('0. 退出');
  
  rl.question('请选择操作: ', async (answer) => {
    switch (answer) {
      case '1':
        await clearCollection();
        showMenu();
        break;
      case '2':
        await addVectorData();
        break;
      case '3':
        await searchVectors();
        break;
      case '0':
        rl.close();
        process.exit(0);
        break;
      default:
        console.log('无效的选择，请重试');
        showMenu();
    }
  });
}

// 启动程序
async function main() {
  try {
    await initCollection();
    showMenu();
  } catch (error) {
    console.error('程序启动失败:', error);
    process.exit(1);
  }
}

main(); 